﻿#include <list>
#include <unordered_set>
#include <unordered_map>
#include <sstream>
#include <memory>

#include "mdp.h"
#include "zsocket.hpp"

#include <absl/time/clock.h>

const static int g_heartbeat_liveness = 10;
const static absl::Duration g_heartbeat_interval = absl::Milliseconds(10000);
const static absl::Duration g_heartbeat_expiry = g_heartbeat_interval * g_heartbeat_liveness;

struct mdworker;
struct mdservice;
using worker_ptr = std::shared_ptr<mdworker>;
using service_ptr = std::shared_ptr<mdservice>;

struct mdworker
{
	std::string m_id_string;
	zmsg m_identity;
	service_ptr m_service;
	absl::Time m_expiry;
};

struct mdservice
{
	mdservice()
	{
		m_workers = 0;
	}

	void remove_worker(worker_ptr worker)
	{
		auto it = std::find(m_waitings.begin(), m_waitings.end(), worker);
		m_waitings.erase(it);
	}

	std::string m_name;
	std::list<zframe> m_requests;
	std::list<worker_ptr> m_waitings;
	int m_workers;
};

class mdbroker
{
public:
	mdbroker(bool verbose) :
		m_socket(m_ctx, zmq::socket_type::router),
		m_verbose(verbose),
		m_heartbeat_at(absl::Now() + g_heartbeat_interval)
	{

	}

public:
	void bind(std::string endpoint)
	{
		m_socket.bind(endpoint);
		info("MDP broker/0.2.0 is active at {}", endpoint);
	}

	void worker_msg(zmsg& sender, zframe& msgs)
	{
		assert(msgs.size() >= 1);

		zmsg command_msg = msgs.pop_front();
		mdpw_command cmd = (mdpw_command)(to_int32(command_msg));
		std::string id_string = to_hex_string(sender);
		bool worker_ready = m_workers.find(id_string) != m_workers.end();
		worker_ptr worker = worker_require(sender);
		switch (cmd)
		{
		case mdpw_command::none:
			break;
		case mdpw_command::ready:
			if (worker_ready) {
				worker_delete(worker, true);
			}
			else if (sender.size() >= 4 &&
				memcpy(sender.data<char>(), "mmi.", 4) == 0)
			{
				worker_delete(worker, true);
			}
			else
			{
				zmsg service_msg = msgs.pop_front();
				std::string service_name = to_string(service_msg);
				worker->m_service = service_require(service_msg);
				worker->m_service->m_waitings.push_back(worker);
				worker->m_service->m_workers++;
				worker_waiting(worker);
			}
			break;
		case mdpw_command::request:
			break;
		case mdpw_command::reply:
			if (worker_ready)
			{
				zmsg client = msgs.pop_front();
				msgs.pop_front();
				msgs.push_front(worker->m_service->m_name);
				msgs.push_front(g_mdpc_client);
				msgs.push_front();
				msgs.push_front(client);
				m_socket.send(msgs);
				worker_waiting(worker);
			}
			else
			{
				worker_delete(worker, true);
			}
			break;
		case mdpw_command::heartbeat:
			if (worker_ready) {
				worker->m_expiry = absl::Now() + g_heartbeat_expiry;
			}
			else {
				worker_delete(worker, true);
			}
			break;
		case mdpw_command::disconnect:
			worker_delete(worker, false);
			break;
		default:
			info("invalid input message");
			msgs.dump();
			break;
		}
	}

	void client_msg(zmsg& sender, zframe& msgs)
	{
		assert(msgs.size() >= 2);

		zmsg service_msg = msgs.pop_front();
		service_ptr service = service_require(service_msg);

		msgs.push_front();
		msgs.push_front(sender);
		if (service_msg.size() >= 4
			&& memcmp(service_msg.data<char>(), "mmi.", 4) == 0)
		{
			std::string return_code;
			if (service_msg == "mmi.service") {
				std::string name = to_string(msgs.back());
				service_ptr service2 = service_require(msgs.back());
				return_code = service2->m_workers > 0 ? "200" : "404";
			}
			else
			{
				return_code = "501";
			}

			msgs.back().rebuild(return_code.data(), return_code.size());
			zmsg client = msgs.pop_front();
			msgs.pop_front();
			msgs.push_front(std::move(service_msg));
			msgs.push_front(g_mdpc_client);
			msgs.push_front();
			msgs.push_front(client);
			m_socket.send(msgs);
		}
		else
		{
			service->m_requests.push_back(std::move(msgs));
			service_dispatch(service);
		}
	}

	void purge()
	{
		std::unordered_set<worker_ptr> to_remove;
		for (auto it = m_workers.begin(); it != m_workers.end(); ++it)
		{
			if (it->second->m_expiry < absl::Now())
			{
				to_remove.insert(it->second);
			}
		}

		for (worker_ptr worker : to_remove)
		{
			if (m_verbose) {
				info("deleting expired worker: {}", worker->m_id_string);
			}
			worker_delete(worker, false);
		}
	}

	void service_dispatch(service_ptr service)
	{
		purge();
		while (!service->m_waitings.empty()
			&& !service->m_requests.empty())
		{
			worker_ptr worker = service->m_waitings.front();
			m_waitings.erase(worker);
			service->m_waitings.pop_front();
			zframe msg = std::move(service->m_requests.front());
			service->m_requests.pop_front();
			worker_send(worker, mdpw_command::request, "", msg.copy());
		}
	}

	void worker_send(worker_ptr worker, mdpw_command cmd, std::string option = "", zframe msg = zframe())
	{
		if (!option.empty()) {
			msg.push_front(option);
		}
		msg.push_front(int32_t(cmd));
		msg.push_front(g_mdpw_worker);
		msg.push_front();
		msg.push_front(worker->m_identity);
		
		if (m_verbose) {
			info("sending {} to worker", g_mdps_commands[int32_t(cmd)]);
			msg.dump();
		}
		m_socket.send(msg);
	}

	service_ptr& service_require(zmsg& id)
	{
		std::string name = to_string(id);
		if (m_services.find(name) == m_services.end())
		{
			if (m_verbose)
			{
				info("added service: {}", name);
			}

			m_services[name] = std::make_shared<mdservice>();
			m_services[name]->m_name = name;
		}
		return m_services[name];
	}

	worker_ptr worker_require(zmsg& id)
	{
		std::string id_string = to_hex_string(id);
		worker_ptr worker = m_workers[id_string];
		if (!worker)
		{
			worker = std::make_shared<mdworker>();
			worker->m_identity.copy(&id);
			worker->m_id_string = id_string;
			m_workers[id_string] = worker;
			if (m_verbose) {
				info("registering new worker: {}", id_string);
			}
		}
		return worker;
	}

	void worker_waiting(worker_ptr worker)
	{
		m_waitings.insert(worker);
		worker->m_service->m_waitings.push_back(worker);
		worker->m_expiry = absl::Now() + g_heartbeat_expiry;

		zframe msg;
		service_dispatch(worker->m_service);
	}

	void worker_delete( worker_ptr worker, bool disconnect)
	{
		if (disconnect) {
			worker_send(worker, mdpw_command::disconnect, "");
		}

		if (worker->m_service) {
			worker->m_service->remove_worker(worker);
			worker->m_service->m_workers--;
		}

		m_waitings.erase(worker);
		m_workers.erase(worker->m_id_string);
	}

	void poll()
	{
		while (true)
		{
			zmq::pollitem_t item = { m_socket, 0, ZMQ_POLLIN, 0 };
			zmq::poll(&item, 1, int(absl::ToInt64Milliseconds(g_heartbeat_interval)));
			if (item.revents & ZMQ_POLLIN)
			{
				zframe msg;
				m_socket.recv(msg);
				if (msg.is_empty())
				{
					break;
				}

				if (m_verbose)
				{
					info("received message:");
					msg.dump();
				}

				zmsg sender = msg.pop_front();
				zmsg empty = msg.pop_front();
				zmsg header = msg.pop_front();
				if (header == g_mdpc_client)
				{
					client_msg(sender, msg);
				}
				else if(header == g_mdpw_worker)
				{
					worker_msg(sender, msg);
				}
				else
				{
					info("invalid message:");
					msg.dump();
				}
			}

			if (absl::Now() > m_heartbeat_at)
			{
				purge();
				for (worker_ptr worker : m_waitings)
				{
					worker_send(worker, mdpw_command::heartbeat);
				}
				m_heartbeat_at = m_heartbeat_at + g_heartbeat_interval;
			}
		}
	}

private:
	zmq::context_t m_ctx;
	zsocket m_socket;
	bool m_verbose;
	std::string m_endpoint;
	std::unordered_map<std::string, worker_ptr> m_workers;
	std::unordered_map<std::string, service_ptr> m_services;
	std::unordered_set<worker_ptr> m_waitings;
	absl::Time m_heartbeat_at;
};

int main(int argc, char* argv[])
{
	bool verbose = (argc > 1 && std::string(argv[1]) == "-v");

	mdbroker broker(verbose);
	broker.bind("tcp://*:5555");
	broker.poll();
}